﻿using System;
using System.IO;
using System.Management.Automation;
using System.Security;
using System.Security.Cryptography;
using System.Text;

namespace CryptoModule
{
    public enum EncryptionAlgorithm
    {
        AES,
        TripleDES,
        Rijndael 

    }

    /// <summary>
    /// Cmdlet to encrypt a file or a string
    /// </summary>
    [Cmdlet("Encrypt","Item")]
    public class EncryptionCmdlet : PSCmdlet
    {
        /// <summary>
        /// The path of the file
        /// </summary>
        [Parameter(Mandatory=true,ParameterSetName="File",ValueFromPipeline=true)]
        [ValidateNotNullOrEmpty]
        public string File
        {
            get;
            set;
        }

        [Parameter(Mandatory=false,ParameterSetName="File")]
        [ValidateNotNullOrEmpty]
        public string Output
        {
            get;
            set;
        }

        /// <summary>
        /// The string to encrypt
        /// </summary>
        [Parameter(Mandatory = true, ParameterSetName = "String",ValueFromPipeline=true)]
        [ValidateNotNullOrEmpty]
        public string Content
        {
            get;
            set;
        }

        /// <summary>
        /// The support Algorithm
        /// </summary>
        [Parameter(ParameterSetName="File")]
        [Parameter(ParameterSetName="String")]
        public EncryptionAlgorithm Algorithm
        {
            get;
            set;
        }


        /// <summary>
        /// The key to encrypt
        /// </summary>
        [Parameter(Mandatory=true,ParameterSetName="File")]
        [Parameter(Mandatory = true, ParameterSetName = "String")]
        [ValidateNotNullOrEmpty]
        public SecureString Key
        {
            get;
            set;
        }

        private SymmetricAlgorithm instance;

        protected override void ProcessRecord()
        {

            instance = Utils.CreateSymmetricAlgorithmInstance(Algorithm);

            if (this.Host.UI != null)
            {
                this.Host.UI.WriteLine(
                    ConsoleColor.Green,
                    Console.BackgroundColor,
                    string.Format("Algorithm {0}. Key size {1} bits, Mode {2}",
                        instance.GetType().Name,
                        instance.KeySize,
                        Enum.GetName(typeof(CipherMode), instance.Mode)));
            }

            instance.GenerateIV();
            Tuple<byte[], byte[]> secret = Utils.ProcessKeyAndIV(Key, instance.KeySize/8);
            instance.Key = secret.Item1;
            instance.GenerateIV();
                                                                                         
            var encryptor = instance.CreateEncryptor();
            if (this.ParameterSetName == "File")
            {
                EncryptFile(encryptor, instance.IV);
            }
            else if (this.ParameterSetName == "String")
            {
                EncryptString(encryptor, instance.IV);
            }
        }

        private void EncryptFile(ICryptoTransform encryptor,byte[] IV)
        {
            if (string.IsNullOrEmpty(Output))
                Output = File + ".enc";

            if (!Path.IsPathRooted(File))
                File = Path.Combine(this.SessionState.Path.CurrentLocation.Path,File);

            if (!Path.IsPathRooted(Output))
                Output = Path.Combine(this.SessionState.Path.CurrentLocation.Path, Output);

            using (FileStream output = new FileStream(Output, FileMode.Create))
            using (FileStream fs = new FileStream(File, FileMode.Open))
            using (CryptoStream cs = new CryptoStream(output, encryptor, CryptoStreamMode.Write))
            using (BinaryWriter sw = new BinaryWriter(cs))
            {
                output.Write(IV, 0, IV.Length);
                Utils.Transform(fs, sw);
                cs.FlushFinalBlock();
                this.WriteObject(Path.GetFullPath(Output));
            }
        }

        private void EncryptString(ICryptoTransform encryptor, byte[] IV)
        {
            using (MemoryStream output = new MemoryStream())
            using (CryptoStream cs = new CryptoStream(output, encryptor, CryptoStreamMode.Write))
            {
                output.Write(IV, 0, IV.Length);
                byte[] data=UnicodeEncoding.Unicode.GetBytes(Content);
                cs.Write(data,0,data.Length);
                cs.FlushFinalBlock();
                this.WriteObject(Convert.ToBase64String(output.ToArray()));
            }
        }
    }

    [Cmdlet("Decrypt", "Item")]
    public class DecryptionCmdlet : PSCmdlet
    {
        [Parameter(Mandatory = true, ParameterSetName = "File")]
        [ValidateNotNullOrEmpty]
        public string File
        {
            get;
            set;
        }

        [Parameter(Mandatory = false, ParameterSetName = "File")]
        [ValidateNotNullOrEmpty]
        public string Output
        {
            get;
            set;
        }

        [Parameter(Mandatory = true, ParameterSetName = "String")]
        [ValidateNotNullOrEmpty]
        public string Content
        {
            get;
            set;
        }

        [Parameter(ParameterSetName = "File")]
        [Parameter(ParameterSetName = "String")]
        public EncryptionAlgorithm Algorithm
        {
            get;
            set;
        }


        [Parameter(Mandatory = true, ParameterSetName = "File")]
        [Parameter(Mandatory = true, ParameterSetName = "String")]
        [ValidateNotNullOrEmpty]
        public SecureString Key
        {
            get;
            set;
        }
        private SymmetricAlgorithm instance;

        protected override void ProcessRecord()
        {

            instance = Utils.CreateSymmetricAlgorithmInstance(Algorithm);
            instance.GenerateIV();
            Tuple<byte[], byte[]> secret = Utils.ProcessKeyAndIV(Key, instance.KeySize/8);
            instance.Key = secret.Item1;

            if (this.ParameterSetName == "File")
            {
                if (!Path.IsPathRooted(File))
                    File = Path.Combine(this.SessionState.Path.CurrentLocation.Path, File);

                instance.IV = Utils.ExtractIVForFile(File, instance.IV.Length);
                var decryptor = instance.CreateDecryptor();
                DecryptFile(decryptor,instance.IV.Length);
            }
            else if (this.ParameterSetName == "String")
            {
                instance.IV = Utils.ExtractIVFromContent(Content, instance.IV.Length);
                var decryptor = instance.CreateDecryptor();
                DecryptString(decryptor,instance.IV.Length);
            }

        }

        private void DecryptFile(ICryptoTransform decryptor, int IVLength)
        {
            if (string.IsNullOrEmpty(Output))
                Output = File + ".dec";

            
            if (!Path.IsPathRooted(Output))
                Output = Path.Combine(this.SessionState.Path.CurrentLocation.Path, Output);

            using (FileStream output = new FileStream(Output, FileMode.Create))
            using (FileStream fs = new FileStream(File, FileMode.Open))
            {
                fs.Seek(IVLength,SeekOrigin.Begin);
                using (CryptoStream cs = new CryptoStream(output, decryptor, CryptoStreamMode.Write))
                using (BinaryWriter sw = new BinaryWriter(cs))
                {
                    Utils.Transform(fs, sw);
                    cs.FlushFinalBlock();
                    this.WriteObject(Path.GetFullPath(Output));
                }
            }
        }

        private void DecryptString(ICryptoTransform decryptor, int IVLength)
        {
            using (MemoryStream output = new MemoryStream())
            using (CryptoStream cs = new CryptoStream(output, decryptor, CryptoStreamMode.Write))
            {
                byte[] data = Convert.FromBase64String(Content);
                cs.Write(data, IVLength, data.Length-IVLength);
                cs.FlushFinalBlock();
                this.WriteObject(UnicodeEncoding.Unicode.GetString(output.ToArray()));
            }
        }
    }
}
